# -*- coding utf-8 -*-
# 作者: SMF
# 时间: 2022.07.16
import copy
import csv

import random

import torch
from torchvision.transforms import ToTensor


def read_csv(path):
    with open(path, 'r', newline='') as csv_file:
        reader = csv.reader(csv_file)
        md_data = []
        md_data += [[float(i) for i in row] for row in reader]
        # md_data_new = ToTensor(md_data)
        return torch.FloatTensor(md_data)


def get_edge_index(matrix):
    edge_index = [[], []]
    for i in range(matrix.size(0)):
        for j in range(matrix.size(1)):
            if matrix[i][j] != 0:
                edge_index[0].append(i)
                edge_index[1].append(j)
    return torch.LongTensor(edge_index)


def read_txt(path):
    with open(path, 'r', newline='') as txt_file:
        reader = txt_file.readlines()
        md_data = []
        md_data += [[float(i) for i in row.split()] for row in reader]
        return torch.FloatTensor(md_data)


def prepare_data(opt):
    dataset = dict()

    # 阅读 miRNA-Disease 关联矩阵
    # dataset['md_p'] = read_csv(opt.data_path + '\\m-d.csv')  # Tensor: (495, 383)
    # dataset['md_true'] = read_csv(opt.data_path + '\\m-d.csv')
    A = read_csv(opt.data_path + '\\m-d.csv')

    zero_index_old = []
    one_index_old = []
    cha_index = {}
    cha_index0 = {}
    cha_index1 = {}
    cha_index2 = {}
    for i in range(A.shape[0]):
        for j in range(A.shape[1]):
            if A[i][j] < 1:
                zero_index_old.append([i, j])
            if A[i][j] >= 1:
                one_index_old.append([i, j])
    random.shuffle(one_index_old)  # (5430, 2)
    # one_index1 = copy.deepcopy(one_index_old)
    # one_index1 = copy.deepcopy(one_index_old)
    random.shuffle(zero_index_old)  # (184155, 2)
    # zero_index1 = copy.deepcopy(zero_index_old)
    # zero_index1 = copy.deepcopy(zero_index_old)

    # 超参数
    length = int(len(one_index_old) * 0.2)  # 20% 的 “1” :1086

    # 五折交叉验证，将矩阵 T 中的“1”划为5份，依次验证
    for i in range(5):
        dataset[i] = dict()
        dataset[i]['md_p'] = copy.deepcopy(A)
        dataset[i]['md_true'] = copy.deepcopy(A)
        # one_index = one_index1
        # zero_index = zero_index1
        one_index = copy.deepcopy(one_index_old)
        zero_index = copy.deepcopy(zero_index_old)
        cha_index1[i] = one_index[i * length:(i + 1) * length]  # 分配入测试集 “1” 的量
        cha_index2[i] = zero_index[i * length:(i + 1) * length]  # 分配入测试集 “0” 的量
        zero_index.extend(one_index[i * length:(i + 1) * length])
        # one_index = one_index[2 * length:len(one_index)]
        del one_index[i * length:(i + 1) * length]

        # 超参数
        length1 = int(len(one_index) * 0.1)  # 分配入验证集的量：434
        cha_index[i] = one_index[i * length1:(i + 1) * length1]  # 分配入验证集 “1” 的量
        cha_index0[i] = zero_index[i * length1:(i + 1) * length1]  # 分配入验证集 “0” 的量
        zero_index.extend(one_index[i * length1:(i + 1) * length1])
        del one_index[i * length1:(i + 1) * length1]
        # one_index = one_index[2 * length:len(one_index)]
        zero_tensor = torch.LongTensor(zero_index)
        one_tensor = torch.LongTensor(one_index)

        # 更新 dataset
        for ind in cha_index[i]:
            if dataset[i]['md_p'][ind[0], ind[1]] != 0:
                dataset[i]['md_p'][ind[0], ind[1]] = 0
            else:
                print("md矩阵有错！")
        for ind1 in cha_index1[i]:
            if dataset[i]['md_p'][ind1[0], ind1[1]] != 0:
                dataset[i]['md_p'][ind1[0], ind1[1]] = 0
            else:
                print("md矩阵有错！")
        dataset[i]['md_p_cha'] = dataset[i]['md_true'] - dataset[i]['md_p']
        dataset[i]['md_true'] = dataset[i]['md_p']

        dataset[i]['md'] = dict()
        dataset[i]['md']['train'] = [one_tensor, zero_tensor]

        # 阅读 Disease-Disease 相似度矩阵
        dd_matrix = read_csv(opt.data_path + '\\d-d.csv')
        dd_edge_index = get_edge_index(dd_matrix)
        dataset[i]['dd'] = {'data': dd_matrix, 'edge_index': dd_edge_index}

        # 阅读 miRNA-miRNA 相似度矩阵
        mm_matrix = read_csv(opt.data_path + '\\m-m.csv')
        mm_edge_index = get_edge_index(mm_matrix)
        dataset[i]['mm'] = {'data': mm_matrix, 'edge_index': mm_edge_index}
    return dataset, cha_index, cha_index0, cha_index1, cha_index2


if __name__ == '__main__':
    # a = read_csv("../dataset/data(383-495)\\m-d.csv")
    pass
